// ExpTrack.C

#include <iostream>
#include <unistd.h>

#include <az_aztec_defs.h>
#include <Epetra_Vector.h>
#include <Epetra_CrsMatrix.h>
#include <NOX.H>
#include <NOX_Epetra.H>
#include <LOCA.H>
#include <LOCA_Epetra.H>

#include "Global.H"
#include "ExpWaveProb.H"              
#include "ExpProbInterface.H"
#include "EigenOp.H"

#define ARGS1    "m n p0 A B contParam stepSize numSteps"
#define NUMARGS1 8
#define ARGS2    "solfile contParam stepSize numSteps"
#define NUMARGS2 4


using namespace std;


int main(int argc, char *argv[])
{
  try
  {
    ExpWaveProb *prob;
    char contParam;
    double stepSize;
    int numSteps;

    if (argc == NUMARGS1+1)
    {
      int i = 1;
      int m = atoi(argv[i++]);
      int n = atoi(argv[i++]);
      double p0 = atof(argv[i++]); if (p0 > 0) { cerr << "Expected negative pressure, exiting" << endl; exit(1); }
      double a = atof(argv[i++]);
      double b = atof(argv[i++]);
      //sprintf(contParam, "%c", toupper(argv[i++][0]));
      contParam = toupper(argv[i++][0]);
      stepSize = atof(argv[i++]);
      numSteps = atoi(argv[i++]);

      prob = new ExpWaveProb(m, n, p0, a, b);
    }
    else if (argc == NUMARGS2+1)
    {
      int i=1;
      char *solfile = argv[i++];
      //sprintf(contParam, "%c", toupper(argv[i++][0]));
      contParam = toupper(argv[i++][0]);
      stepSize = atof(argv[i++]);
      numSteps = atoi(argv[i++]);

      prob = new ExpWaveProb(solfile);
    }
    else { fprintf(stderr, "usage:\n  %s %s\n  %1$s %3$s\n", argv[0], ARGS1, ARGS2); exit(1); }

    double contParamVal;
    switch (contParam) {
    case 'A':
      contParamVal = prob->getA();
      break;
    case 'B':
      contParamVal = prob->getB();
      break;
    case 'Q':
      contParamVal = prob->getQ();
      break;
    default:
      cout << "Bad continuation parameter " << contParam << endl;
      exit(1);
    }

    char contParamS[2];
    sprintf(contParamS,"%c",contParam);


    // Get the solution vector from the problem
    Epetra_Vector& soln = prob->getInitial();

    // Begin LOCA Solver ************************************

    Teuchos::RCP<Teuchos::ParameterList> paramList = 
      Teuchos::rcp(new Teuchos::ParameterList);

    // Create LOCA sublist
    Teuchos::ParameterList& locaParamsList = paramList->sublist("LOCA");

    // Create the stepper sublist and set the stepper parameters
    Teuchos::ParameterList& locaStepperList = locaParamsList.sublist("Stepper");
    locaStepperList.set("Continuation Method", "Arc Length"); // NOTE used to be "Householder Arc Length"
    locaStepperList.set("Continuation Parameter", contParamS);
    locaStepperList.set("Initial Value", contParamVal);
    locaStepperList.set("Max Value", 5000.0);
    locaStepperList.set("Min Value", 0.0);
    locaStepperList.set("Max Steps", numSteps);
    locaStepperList.set("Max Nonlinear Iterations", 20);
    locaStepperList.set("Enable Arc Length Scaling", true);
    locaStepperList.set("Goal Arc Length Parameter Contribution", 0.5);
    locaStepperList.set("Max Arc Length Parameter Contribution", 0.7);
    locaStepperList.set("Initial Scale Factor", 1.0);
    locaStepperList.set("Min Scale Factor", 1.0e-8);
    locaStepperList.set("Enable Tangent Factor Step Size Scaling", true);
    locaStepperList.set("Min Tangent Factor", -1.8);
    locaStepperList.set("Tangent Factor Exponent", 1.0);

    // Create bifurcation sublist
    Teuchos::ParameterList& bifurcationList = 
      locaParamsList.sublist("Bifurcation");
    bifurcationList.set("Method", "None");

    // Create Anasazi Eigensolver sublist (needs --with-loca-anasazi)
    locaStepperList.set("Compute Eigenvalues", true);
    Teuchos::ParameterList& eList = 
      locaStepperList.sublist("Eigensolver");
    eList.set("Method","Anasazi");
    eList.set("Block Size", 1);   // Size of blocks
    eList.set("Num Blocks", 30);  // Size of Arnoldi factorization
    eList.set("Num Eigenvalues", 4);  // Number of eigenvalues
    eList.set("Convergence Tolerance", 2.0e-5);  // Tolerance
    eList.set("Step Size", 1);  // How often to check convergence
    eList.set("Maximum Restarts", 2);
    eList.set("Sorting Order", "SR");  // changing this to SM will break things
    eList.set("Verbosity",
	      Anasazi::Errors + 
	      Anasazi::Warnings +
	      Anasazi::FinalSummary);        // Verbosity
    eList.set("Save Eigen Data Method", "User-Defined");
    eList.set("User-Defined Save Eigen Data Name", "CustomEigenOp");
    Teuchos::RCP<LOCA::SaveEigenData::AbstractStrategy> sedm =
      Teuchos::rcp(new EigenOp);
    eList.set("CustomEigenOp", sedm);

    // Create predictor sublist
    Teuchos::ParameterList& predictorList = 
      locaParamsList.sublist("Predictor");
    //predictorList.set("Method", "Tangent");
    predictorList.set("Method", "Secant"); 
 
    Teuchos::ParameterList& firstStepPredictor= predictorList.sublist("First Step Predictor");
    firstStepPredictor.set("Method", "Random");
    firstStepPredictor.set("Epsilon", 1.0e-3);

    Teuchos::ParameterList& lastStepPredictor = predictorList.sublist("Last Step Predictor");
    lastStepPredictor.set("Method", "Random");
    lastStepPredictor.set("Epsilon", 1.0e-3);
    // Create step size sublist
    Teuchos::ParameterList& stepSizeList = locaParamsList.sublist("Step Size");
    stepSizeList.set("Initial Step Size", stepSize);
    stepSizeList.set("Method", "Adaptive");
    stepSizeList.set("Min Step Size", 1.0e-5);
    stepSizeList.set("Max Step Size", 10.0);
    stepSizeList.set("Aggressiveness", 0.5);
    stepSizeList.set("Failed Step Reduction Factor", 0.5);

    // Set the LOCA Utilities
    Teuchos::ParameterList& locaUtilsList = locaParamsList.sublist("Utilities");
    //locaUtilsList.set("MyPID", MyPID);
// TODO: replace the following:
//    locaUtilsList.set("Output Information", 
//        LOCA::Utils::Warning +
//        LOCA::Utils::StepperIteration +
//        LOCA::Utils::StepperDetails +
//        LOCA::Utils::Solver +
//        LOCA::Utils::SolverDetails +
//        LOCA::Utils::Parameters);

    // Create the "Solver" parameters sublist to be used with NOX Solvers
    Teuchos::ParameterList& nlParams = paramList->sublist("NOX");
    nlParams.set("Nonlinear Solver", "Line Search Based"); 

    // Create the NOX printing parameter list
    Teuchos::ParameterList& nlPrintParams = nlParams.sublist("Printing");
    //nlPrintParams.set("MyPID", MyPID); 
    nlPrintParams.set("Output Information", NOX::Utils::Warning);
    nlPrintParams.set("Output Information",
        NOX::Utils::OuterIteration + 
        NOX::Utils::OuterIterationStatusTest + 
        NOX::Utils::InnerIteration +
        NOX::Utils::Parameters + 
        NOX::Utils::Details + 
        NOX::Utils::Warning);

    // Create the "Line Search" sublist for the "Line Search Based" solver
    Teuchos::ParameterList& searchParams = nlParams.sublist("Line Search");
    searchParams.set("Method", "Backtrack");
    Teuchos::ParameterList& backtrackParams = searchParams.sublist("Backtrack");
    backtrackParams.set("Max Iters", 10);
    backtrackParams.set("Default Step", 1.0000);
    backtrackParams.set("Recovery Step", 0.0001);
    backtrackParams.set("Minimum Step", 0.0001);

    // Create the "Direction" sublist for the "Line Search Based" solver
    Teuchos::ParameterList& dirParams = nlParams.sublist("Direction");
    Teuchos::ParameterList& newParams = dirParams.sublist("Newton");
    dirParams.set("Method", "Newton");
    newParams.set("Forcing Term Method", "Constant");

    // Create the "Linear Solver" sublist for the "Direction" sublist
    Teuchos::ParameterList& lsParams = newParams.sublist("Linear Solver");
    lsParams.set("Aztec Solver", "GMRES");  
    lsParams.set("Max Iterations", 80);  
    lsParams.set("Tolerance", 1e-6);
    lsParams.set("Output Frequency", 40);
    //lsParams.set("Preconditioner", "AztecOO");
    //lsParams.set("Aztec Preconditioner", "ilut"); 
    lsParams.set("Preconditioner", "New Ifpack");
    // Preconditioner used to be AztecOO, but was causing recompute error

    // Create and initialize the parameter vector
    LOCA::ParameterVector pVector;
    pVector.addParameter(contParamS, contParamVal);
    
    // Create the interface between the test problem and the nonlinear solver
    // This is created by the user using inheritance of the abstract base class:
    // NLS_PetraGroupInterface
    Teuchos::RCP<ExpProbInterface> interface = 
      Teuchos::rcp (new ExpProbInterface(*prob, contParamS));
    Teuchos::RCP<LOCA::Epetra::Interface::Required> iReq = interface;
    Teuchos::RCP<NOX::Epetra::Interface::Jacobian> iJac = interface;

    // Create the Epetra_RowMatrix for the Jacobian/Preconditioner
    Teuchos::RCP<Epetra_RowMatrix> Amat = 
      Teuchos::rcp(&prob->getJacobian(),false);

    // Create the linear systems
    Teuchos::RCP<NOX::Epetra::LinearSystemAztecOO> linsys = 
      Teuchos::rcp(new NOX::Epetra::LinearSystemAztecOO(nlPrintParams, lsParams,
						      iReq, iJac, Amat, soln));

    // Create the loca vector
    NOX::Epetra::Vector locaSoln(soln);

  // Create Epetra factory
  Teuchos::RCP<LOCA::Abstract::Factory> epetraFactory =
    Teuchos::rcp(new LOCA::Epetra::Factory);

  // Create global data object
  Teuchos::RCP<LOCA::GlobalData> globalData = 
    LOCA::createGlobalData(paramList, epetraFactory);

  // Create the Group
  Teuchos::RCP<LOCA::Epetra::Group> grp = 
    Teuchos::rcp(new LOCA::Epetra::Group(globalData, nlPrintParams, 
					 iReq, locaSoln, linsys,
					 pVector));

  grp->computeF();

  // Create the Solver convergence test
  Teuchos::RCP<NOX::StatusTest::NormF> wrms = 
    Teuchos::rcp(new NOX::StatusTest::NormF(1.0e-6));
  Teuchos::RCP<NOX::StatusTest::MaxIters> maxiters = 
    Teuchos::rcp(new NOX::StatusTest::MaxIters(locaStepperList.get("Max Iters", 6)));
  Teuchos::RCP<NOX::StatusTest::Combo> combo = 
    Teuchos::rcp(new NOX::StatusTest::Combo(NOX::StatusTest::Combo::OR));
  combo->addStatusTest(wrms);
  combo->addStatusTest(maxiters);

  // Create the stepper  
  LOCA::Stepper stepper(globalData, grp, combo, paramList);

  // Perform continuation run
  LOCA::Abstract::Iterator::IteratorStatus status = stepper.run();

  if (status == LOCA::Abstract::Iterator::Finished) 
    globalData->locaUtils->out() << "\nContinuation finished." << endl;
  else {
    if (globalData->locaUtils->isPrintType(NOX::Utils::Error))
      globalData->locaUtils->out() 
	<< "Stepper failed to converge!" << std::endl;
    }

  // Output the parameter list
  if (globalData->locaUtils->isPrintType(NOX::Utils::StepperParameters)) {
    globalData->locaUtils->out() 
      << std::endl << "Final Parameters" << std::endl
      << "****************" << std::endl;
    stepper.getList()->print(globalData->locaUtils->out());
    globalData->locaUtils->out() << std::endl;
  }

  LOCA::destroyGlobalData(globalData);
    delete prob;
  }
  catch (char *s) { cout << s << endl; }
  //catch (...) { cout << "Caught unknown exception" << endl; }
}
